import os
import time
import unittest
from unittest import skipIf
from unittest import skipUnless
from unittest import TestCase

from packaging import version

import mongomock
import mongomock.gridfs
from mongomock import helpers


try:
    import gridfs
    from gridfs import errors

    _HAVE_GRIDFS = True
except ImportError:
    _HAVE_GRIDFS = False


try:
    import pymongo
    from bson.objectid import ObjectId
    from pymongo import MongoClient as PymongoClient
except ImportError:
    ...


@skipUnless(helpers.HAVE_PYMONGO, 'pymongo not installed')
@skipUnless(_HAVE_GRIDFS and hasattr(gridfs.__builtins__, 'copy'), 'gridfs not installed')
@skipIf(os.getenv('NO_LOCAL_MONGO'), 'No local Mongo server running')
class GridFsTest(TestCase):
    @classmethod
    def setUpClass(cls):
        mongomock.gridfs.enable_gridfs_integration()

    def setUp(self):
        super().setUp()
        self.fake_conn = mongomock.MongoClient()
        self.mongo_conn = self._connect_to_local_mongodb()
        self.db_name = 'mongomock___testing_db'

        self.mongo_conn[self.db_name]['fs']['files'].drop()
        self.mongo_conn[self.db_name]['fs']['chunks'].drop()

        self.real_gridfs = gridfs.GridFS(self.mongo_conn[self.db_name])
        self.fake_gridfs = gridfs.GridFS(self.fake_conn[self.db_name])

    def tearDown(self):
        super().setUp()
        self.mongo_conn.close()
        self.fake_conn.close()

    def test__put_get_small(self):
        before = time.time()
        fid = self.fake_gridfs.put(GenFile(50))
        rid = self.real_gridfs.put(GenFile(50))
        after = time.time()
        ffile = self.fake_gridfs.get(fid)
        rfile = self.real_gridfs.get(rid)
        self.assertEqual(ffile.read(), rfile.read())
        fake_doc = self.get_fake_file(fid)
        mongo_doc = self.get_mongo_file(rid)
        self.assertSameFile(mongo_doc, fake_doc, max_delta_seconds=after - before + 1)

    def test__put_get_big(self):
        # 500k files are bigger than doc size limit
        before = time.time()
        fid = self.fake_gridfs.put(GenFile(500000, 10))
        rid = self.real_gridfs.put(GenFile(500000, 10))
        after = time.time()
        ffile = self.fake_gridfs.get(fid)
        rfile = self.real_gridfs.get(rid)
        self.assertEqual(ffile.read(), rfile.read())
        fake_doc = self.get_fake_file(fid)
        mongo_doc = self.get_mongo_file(rid)
        self.assertSameFile(mongo_doc, fake_doc, max_delta_seconds=after - before + 1)

    def test__delete_exists_small(self):
        fid = self.fake_gridfs.put(GenFile(50))
        self.assertTrue(self.get_fake_file(fid) is not None)
        self.assertTrue(self.fake_gridfs.exists(fid))
        self.fake_gridfs.delete(fid)
        self.assertFalse(self.fake_gridfs.exists(fid))
        self.assertFalse(self.get_fake_file(fid) is not None)
        # All the chunks got removed
        self.assertEqual(0, self.fake_conn[self.db_name].fs.chunks.count_documents({}))

    def test__delete_exists_big(self):
        fid = self.fake_gridfs.put(GenFile(500000))
        self.assertTrue(self.get_fake_file(fid) is not None)
        self.assertTrue(self.fake_gridfs.exists(fid))
        self.fake_gridfs.delete(fid)
        self.assertFalse(self.fake_gridfs.exists(fid))
        self.assertFalse(self.get_fake_file(fid) is not None)
        # All the chunks got removed
        self.assertEqual(0, self.fake_conn[self.db_name].fs.chunks.count_documents({}))

    def test__delete_no_file(self):
        # Just making sure we don't crash
        self.fake_gridfs.delete(ObjectId())

    def test__list_files(self):
        fids = [
            self.fake_gridfs.put(GenFile(50, 9), filename='one'),
            self.fake_gridfs.put(GenFile(62, 5), filename='two'),
            self.fake_gridfs.put(GenFile(654, 1), filename='three'),
            self.fake_gridfs.put(GenFile(5), filename='four'),
        ]
        names = ['one', 'two', 'three', 'four']
        names_no_two = [x for x in names if x != 'two']
        for x in self.fake_gridfs.list():
            self.assertIn(x, names)

        self.fake_gridfs.delete(fids[1])

        for x in self.fake_gridfs.list():
            self.assertIn(x, names_no_two)

        three_file = self.get_fake_file(fids[2])
        self.assertEqual('three', three_file['filename'])
        self.assertEqual(654, three_file['length'])
        self.fake_gridfs.delete(fids[0])
        self.fake_gridfs.delete(fids[2])
        self.fake_gridfs.delete(fids[3])
        self.assertEqual(0, len(self.fake_gridfs.list()))

    def test__find_files(self):
        file_ids = []
        for name, data in [
            ('a', GenFile(50, 9)),
            ('b', GenFile(62, 5)),
            ('b', GenFile(654, 1)),
            ('a', GenFile(5)),
        ]:
            time.sleep(0.001)
            file_ids.append(self.fake_gridfs.put(data, filename=name))

        c = self.fake_gridfs.find({'filename': 'a'}).sort('uploadDate', -1)
        file3 = c.next()
        file0 = c.next()
        self.assertFalse(c.alive)
        self.assertNotEqual(file3.uploadDate, file0.uploadDate)

        self.assertEqual(file_ids[3], file3._id)
        self.assertEqual(file_ids[0], file0._id)

    def test__put_exists(self):
        self.fake_gridfs.put(GenFile(1), _id='12345')
        with self.assertRaises(errors.FileExists):
            self.fake_gridfs.put(GenFile(2, 3), _id='12345')

    def assertSameFile(self, real, fake, max_delta_seconds=1):
        # https://pymongo.readthedocs.io/en/stable/migrate-to-pymongo4.html#disable-md5-parameter-is-removed
        if version.parse('4.0') > helpers.PYMONGO_VERSION:
            self.assertEqual(real['md5'], fake['md5'])

        self.assertEqual(real['length'], fake['length'])
        self.assertEqual(real['chunkSize'], fake['chunkSize'])
        self.assertLessEqual(
            abs(real['uploadDate'] - fake['uploadDate']).seconds,
            max_delta_seconds,
            msg='real: {}, fake: {}'.format(real['uploadDate'], fake['uploadDate']),
        )

    def get_mongo_file(self, i):
        return self.mongo_conn[self.db_name]['fs']['files'].find_one({'_id': i})

    def get_fake_file(self, i):
        return self.fake_conn[self.db_name]['fs']['files'].find_one({'_id': i})

    def _connect_to_local_mongodb(self, num_retries=60):
        """Performs retries on connection refused errors (for travis-ci builds)"""
        for retry in range(num_retries):
            if retry > 0:
                time.sleep(0.5)
            try:
                return PymongoClient(
                    host=os.environ.get('TEST_MONGO_HOST', 'localhost'), maxPoolSize=1
                )
            except pymongo.errors.ConnectionFailure as e:
                if retry == num_retries - 1:
                    raise
                if 'connection refused' not in e.message.lower():
                    raise


class GenFile:
    def __init__(self, length, value=0, do_encode=True):
        self.gen = self._gen_data(length, value)
        self.do_encode = do_encode

    def _gen_data(self, length, value):
        while length:
            length -= 1
            yield value

    def _maybe_encode(self, s):
        if self.do_encode and isinstance(s, str):
            return s.encode('UTF-8')
        return s

    def read(self, num_bytes=-1):
        s = ''
        bytes_left = -1 if num_bytes <= 0 else num_bytes
        while True:
            n = next(self.gen, None)
            if n is None:
                return self._maybe_encode(s)
            s += chr(n)
            bytes_left -= 1
            if bytes_left == 0:
                return self._maybe_encode(s)


if __name__ == '__main__':
    unittest.main()
